"""
Author        yongfa
Date          2022-10-19 16:14:38
LastEditTime  2022-10-21 16:32:21
LastEditors   yongfa
Description   pymongo 封装类
"""


import pymongo
import queue
from pymongo import UpdateOne, UpdateMany, ReplaceOne


class MongoConnectionPool(object):

    def __init__(self, params):
        if 'mongo_url' not in params:
            raise Exception('未指定数据连接地址: mongo_url.')

        if 'db' not in params:
            raise Exception('未指定数据库')

        self.mongo_url = params['mongo_url']
        self.db = params['db']
        self.pool_size = 5 if 'pool_size' not in params else params['pool_size']

        self.conn_queue = queue.Queue(maxsize=self.pool_size)
        for _ in range(self.pool_size):
            self.conn_queue.put(self.__create_new_conn())

    def __create_new_conn(self):

        connect = pymongo.MongoClient(self.mongo_url)[self.db]
        return connect

    def get_conn(self):
        conn = self.conn_queue.get()
        if conn is None:
            conn = self.__create_new_conn()
        return conn

    def put_conn(self, conn):
        self.conn_queue.put(conn)


class MongoHelper:
    def __init__(self, params: dict = None, pool: MongoConnectionPool = None):
        self.pool = None
        if pool:
            self.pool = pool

        else:
            if 'mongo_url' not in params:
                raise Exception('未指定数据连接地址: mongo_url.')
            if 'db' not in params:
                raise Exception('未指定数据库')

            self.mongo_url = params['mongo_url']
            self.db = params['db']

            self.client = pymongo.MongoClient(self.mongo_url)[self.db]

    def insert_init(self, collection, size=500):
        self.collection = collection
        self.bucket_ = []
        self.bulk_size = size
        self.bulk_total = 0

    @property
    def bucket(self):
        if self.bucket_:
            if len(self.bucket_) == self.bulk_size or self.bucket_[-1] is None:
                self.bucket_.pop()
                data = []
                for i in self.bucket_:
                    data.append(
                        UpdateOne(
                            {'_id': i[0]},
                            {'$set': i[1]},
                            upsert=True
                        )
                    )
                if data:
                    self.bulk_write(self.collection, data)
                    self.bulk_total += len(data)
                    self.bucket_ = []

        return self.bucket_

    def insert_one(self, collection, doc):
        """单条插入"""
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].insert_one(doc).inserted_id
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].insert_one(doc).inserted_id

    def insert_many(self, collection, docs=None, update=True):
        """多条插入"""
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].insert_many(docs).inserted_ids
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].insert_many(docs).inserted_ids

    def delete_one(self, collection, dsl):
        """单条插入"""
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].delete_one(dsl)
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].delete_one(dsl)

    def delete_many(self, collection, dsl):
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].delete_many(dsl)
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].delete_many(dsl)

    def update_one(self, collection, dsl, col, upsert=False):
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].update_one(dsl, col, upsert=upsert)
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].update_one(dsl, col, upsert=upsert)

    def update_many(self, collection: str, dsl: dict, col: dict):
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].update_many(dsl, col)
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].update_many(dsl, col)

    def find_one(self, collection: str, dsl: dict = None, ret: dict = None,):
        if self.pool:
            client = self.pool.get_conn()

            if ret:
                data = client[collection].find_one(dsl, ret)
            else:
                data = client[collection].find_one(dsl)

            self.pool.put_conn(client)

            return data

        else:
            if ret:
                return self.client[collection].find_one(dsl, ret)
            else:
                return self.client[collection].find_one(dsl)

    def find(self, collection, *args, **kwargs):
        if self.pool:
            client = self.pool.get_conn()
        else:
            client = self.client

        res = client[collection].find(*args, **kwargs)

        if 'sort' in kwargs:
            res = res.sort(kwargs['sort'])

        if 'skip' in kwargs:
            res = res.skip(kwargs['skip'])

        if 'limit' in kwargs:
            res = res.limit(kwargs['limit'])

        if self.pool:
            self.pool.put_conn(client)

        data = [i for i in res]

        lastID = data[-1].get('_id') if data else None

        return data, lastID

    def bulk_write(self, collection, requests, ordered=True):
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].bulk_write(requests, ordered=ordered)
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].bulk_write(requests, ordered=ordered)

    def drop(self, collection):
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].drop()
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].drop()

    def list_collections(self):
        if self.pool:
            client = self.pool.get_conn()
            data = client.list_collections()
            self.pool.put_conn(client)
            return data
        else:
            return self.client.list_collections()

    def create_index(self, collection, index):
        if self.pool:
            client = self.pool.get_conn()
            data = client[collection].create_index(index)
            self.pool.put_conn(client)
            return data
        else:
            return self.client[collection].create_index(index)

    def close(self):
        if self.pool:
            client = self.pool.get_conn()
            return client.close()
        else:
            return self.client.close()


if __name__ == '__main__':
    mongo_params = {
        'pool_size': 5,
        'db': 'runoobdb',
        'mongo_url': 'mongodb://admin:admin321%3F@localhost:27017'
    }

    mongoClient = MongoHelper(pool=MongoConnectionPool(mongo_params))
    # mongoClient = MongoHelper(params=mongo_params)
    import time
    col_list = mongoClient.list_collections()
    print(col_list)

    # v = mongoClient.insert_one('test', {'id': 'admin', 'content': 'admin321%3F'})
    # print(v)

    # 单条插入
    s_time = time.time()
    # for i in range(100000):
    #     mongoClient.insert_one('test', {'id': 'admin{}'.format(str(i)), 'content': 'content{}'.format(
    #         str(i)), 'title': 'sadfsadfljsdlfk;jdslfkjsa;lkdfjskla;fjas;lkfjd'})
    # print('cost: ', time.time() - s_time)

    # 批量插入
    # s_time = time.time()
    # mongoClient.insert_many('test_many', [{'id': 'admin{}'.format(
    #     str(i)), 'content': 'content{}'.format(str(i)), "title": "abcdefg"} for i in range(1000000)])
    # print('cost: ', time.time() - s_time)
    # data = [UpdateOne({"id": "admin"}, {"$set": {"content": "bulk_write"}}, upsert=True)]

    # data = [UpdateMany({"id": "admin"}, {"$set": {"content": "bulk_write", "title": "hello"}}, upsert=True)]
    # data = [ReplaceOne({"id": "admin"}, {"content": "replaceOne"}, upsert=True)]
    # m = mongoClient.bulk_write('test', data)

    # data = [UpdateOne({"id": i}, {"$set": {"content": "bulk_write", "title": "hello"}}, upsert=True)
    #         for i in range(100000)]

    # mongoClient.bulk_write('upsert', data)
    # print('cost: ', time.time() - s_time)

    # 单一条件
    # data = mongoClient.find('test', {"id": "admin"}).explain("executionStats")
    # print(data)
    # import time
    # s_time = time.time()
    # data = mongoClient.find('test', {"id": {'$in': ["admin", "admin3"]}}, {"content": 1, "id": 1})
    # print('cost: ', time.time() - s_time)
    # print(data)

    # 分页查询

    # pages = 1
    # page_size = 10

    # for i in range(0, pages):

    #     kwargs = {'skip': i * page_size, 'limit': page_size, "sort": [('id', -1)]}
    #     result, lastID = mongoClient.find('upsert', {}, **kwargs)
    #     print(result)
    #     print('lastID:', lastID)

    # print('cost: ', time.time() - s_time)

    # 创建索引
    data = mongoClient.create_index('test_many', [('id', 1)])
    print('cost: ', time.time() - s_time)
    print(data)
